"""
可行性研究 v2：基于重构后的模块实现跳层推理实验
"""
from llm.skipable_model import SkipableModel
from env.llm.model_loader import ModelLoader


def main():
    """
    执行跳层推理实验 - 基于重构后的架构
    """
    # --- 1. 配置 ---
    MODEL_ID = "meta-llama/Llama-2-7b-chat-hf"
    CACHE_DIR = "./hf_cache/"
    
    # --- 2. 加载模型 ---
    loader = ModelLoader(cache_dir=CACHE_DIR)
    model, tokenizer = loader.load_model(MODEL_ID)
    
    # 创建可跳层模型
    skipable_model = SkipableModel(model, tokenizer)
    
    # --- 3. 准备测试用例 ---
    test_prompts = [
        # 逻辑推理问题
        "Which name is longer? Bob has 3 letters, Alice has 5 letters. Therefore, Alice is",
        
        # # 数学比较问题  
        # "Compare 8.11 and 8.9: Since 8.11 has more digits after decimal but 8.9 = 8.90, we see 8.9 is actually",
        
        # # 算术问题
        # "What is 15 + 27? Let me calculate step by step: 15 + 27 = 10 + 5 + 20 + 7 = 30 + 12 = ",
        
        # # 几何问题
        # "A circle has radius 5. What is its area? Using the formula A = πr², we get A = π × 5² = π × ",
        
        # # 代数问题
        # "Solve for x: 2x + 6 = 14. First subtract 6 from both sides: 2x = 8. Then divide by 2: x = ",
        
        # # 概率问题
        # "A fair coin is flipped twice. What is the probability of getting at least one head? The probability is 1 - P(no heads) = 1 - (1/2)² = 1 - ",
        
        # # 编程逻辑问题
        # 'In Python, what does this code print?\nfor i in range(3):\n    print(i)\nThe output will be: 0, 1, ',
        
        # # 物理问题
        # "An object falls from height h. Using v² = u² + 2as with u=0, a=g, s=h, the final velocity is v = √(2gh). If g=10 m/s² and h=5m, then v = √(2×10×5) = √(",
        
        # # 化学问题
        # "Balance this equation: C + O₂ → CO₂. We need 1 carbon atom and 2 oxygen atoms on each side, so the balanced equation is: C + ",
        
        # # 历史问题
        # "World War II ended in 1945. It started in 1939 when Germany invaded Poland. Therefore, the war lasted approximately",
        
        # # 文学分析
        # "In Shakespeare's Romeo and Juliet, the main theme is the conflict between love and hate. The two feuding families are the Montagues and the",
        
        # # 生物学问题
        # "DNA is made of four bases: A, T, G, and C. In DNA replication, A pairs with T, and G pairs with C. So if one strand is ATGC, the complementary strand is",
        
        # # 计算机科学问题
        # "What is the time complexity of binary search? Since we divide the search space in half each time, the complexity is O(log n). If we have 1000 elements, we need at most log₂(1000) ≈ ",
        
        # # 经济学问题
        # "If inflation is 3% per year, what will $100 be worth after 2 years? Using compound interest: Value = 100 × (1.03)² = 100 × "
    ]
    
    # --- 4. 执行实验 ---
    num_layers = skipable_model.num_layers
    all_results = []
    
    for prompt_idx, prompt in enumerate(test_prompts):
        print(f"\n" + "="*50)
        print(f" 测试例子 {prompt_idx + 1}: '{prompt}' ")
        print("="*50)
        
        # 完整推理 (基准)
        print("\n执行完整推理 (基准)")
        print("-" * 30)
        
        baseline_token = skipable_model.full_inference(prompt)
        print(f"提示: '{prompt}'")
        print(f"基准模型预测的下一个词元: '{baseline_token}'")
        
        # 逐层跳过进行推理
        print("\n执行跳层推理实验")
        print("-" * 30)
        
        # 记录当前prompt的结果
        prompt_results = {
            'prompt': prompt,
            'baseline_token': baseline_token,
            'skip_results': [],
            'error_layers': []
        }
        
        for layer_to_skip in range(num_layers):
            skipped_token = skipable_model.skip_single_layer(prompt, layer_to_skip)
            is_same = (skipped_token == baseline_token)
            
            # 记录结果
            prompt_results['skip_results'].append({
                'layer_skipped': layer_to_skip,
                'predicted_token': skipped_token,
                'is_correct': is_same
            })
            
            if not is_same:
                prompt_results['error_layers'].append(layer_to_skip)
            
            print(f"跳过第 {layer_to_skip:02d} 层 | 预测词元: '{skipped_token}' | "
                  f"与基准一致: {'✅' if is_same else '❌'}")
        
        all_results.append(prompt_results)
    
    # --- 5. 统计分析 ---
    print("\n" + "="*60)
    print(" 综合统计分析 ")
    print("="*60)
    
    total_tests = len(test_prompts) * num_layers
    total_errors = 0
    all_error_layers = []
    
    for result in all_results:
        error_count = len(result['error_layers'])
        total_errors += error_count
        all_error_layers.extend(result['error_layers'])
        
        print(f"\n提示: '{result['prompt']}'")
        print(f"  基准预测: '{result['baseline_token']}'")
        print(f"  错误层数: {error_count}/{num_layers}")
        print(f"  错误率: {error_count/num_layers*100:.1f}%")
        if result['error_layers']:
            print(f"  错误层序号: {result['error_layers']}")
    
    # 计算总体统计
    overall_error_rate = total_errors / total_tests * 100
    
    print(f"\n" + "-"*40)
    print(" 总体统计 ")
    print("-"*40)
    print(f"总测试次数: {total_tests}")
    print(f"总错误次数: {total_errors}")
    print(f"总体错误率: {overall_error_rate:.2f}%")
    
    if all_error_layers:
        mean_error_layer = sum(all_error_layers) / len(all_error_layers)
        print(f"错误层序号均值: {mean_error_layer:.2f}")
        
        # 计算错误层的分布
        layer_error_count = {}
        for layer in all_error_layers:
            layer_error_count[layer] = layer_error_count.get(layer, 0) + 1
        
        print(f"\n错误层分布:")
        for layer in sorted(layer_error_count.keys()):
            count = layer_error_count[layer]
            percentage = count / len(all_error_layers) * 100
            print(f"  第{layer:02d}层: {count}次 ({percentage:.1f}%)")
    else:
        print("错误层序号均值: 无错误")
        print("所有跳层测试都与基准预测一致！")


if __name__ == "__main__":
    main()
